import random
import gym
import torch.optim
import numpy as np
import torch.nn.functional as F
from matplotlib import pyplot as plt
from tqdm import tqdm
import utils
from Net import Qnet, ConvolutionQnet, VAnet, ConvolutionVAnet
from ReplayBuffer import ReplayBuffer


class DQN:
    """DQN算法 （包含两个Q网络） 包括Double DQN"""

    def __init__(self, state_dim, hidden_dim, action_dim, learning_rate, gamma, epsilon, target_update, device,
                 isConv=False, input_shape=None, dqn_type="VanillaDQN"):
        self.action_dim = action_dim

        if isConv:  # 图像作为输入时
            if dqn_type == "DuelingDQN":
                # 用于学习真正Q函数的Q网络
                self.q_net = ConvolutionVAnet(input_shape, action_dim, in_channels=input_shape[0]).to(device)
                # 目标网络，用于稳定更新
                self.target_q_net = ConvolutionVAnet(input_shape, action_dim, in_channels=input_shape[0]).to(device)
            else:
                self.q_net = ConvolutionQnet(input_shape, action_dim, in_channels=input_shape[0]).to(device)
                self.target_q_net = ConvolutionQnet(input_shape, action_dim, in_channels=input_shape[0]).to(device)
        else:
            if dqn_type == "DuelingDQN":
                self.q_net = VAnet(state_dim, hidden_dim, action_dim).to(device)
                self.target_q_net = VAnet(state_dim, hidden_dim, action_dim).to(device)
            else:
                self.q_net = Qnet(state_dim, hidden_dim, action_dim).to(device)
                self.target_q_net = Qnet(state_dim, hidden_dim, action_dim).to(device)



        # 优化器
        self.optimizer = torch.optim.Adam(self.q_net.parameters(), lr=learning_rate)
        # 折扣因子
        self.gamma = gamma
        # epsilon-贪婪策略
        self.epsilon = epsilon
        # 目标网络参数更新频率
        self.target_update = target_update
        # 计数器，记录更新次数
        self.count = 0
        self.device = device
        # DQN类型
        self.dqn_type = dqn_type

    def take_action(self, state):
        if np.random.random() < self.epsilon:
            action = np.random.randint(self.action_dim)
        else:
            state = torch.tensor(np.array([state]), dtype=torch.float).to(self.device)
            # 获取输出值(Q值)最大的那个action
            action = self.q_net(state).argmax().item()
        return action

    def numpy_to_tensor_gpu(self, data):
        return torch.tensor(data).to(self.device)

    def max_q_values(self, state):
        """Double DQN 中使用训练的q_net来选取动作"""
        state = torch.tensor([state], dtype=torch.float).to(self.device)
        return self.q_net(state).max().item()

    def update(self, transition_dict):
        """
        更新方式Q网络
        :param transition_dict: 一批轨迹数据
        """
        states = torch.tensor(transition_dict['states'], dtype=torch.float).to(self.device)
        actions = torch.tensor(transition_dict['actions'], dtype=torch.int64).view(-1, 1).to(self.device)
        rewards = torch.tensor(transition_dict['rewards'], dtype=torch.float).view(-1, 1).to(self.device)
        next_states = torch.tensor(transition_dict['next_states'], dtype=torch.float).to(self.device)
        dones = torch.tensor(transition_dict['dones'], dtype=torch.float).view(-1, 1).to(self.device)

        # 获取对应(s,a)的Q(s,a)
        q_values = self.q_net(states)
        q_values = q_values.gather(1, actions)
        if self.dqn_type == "DoubleDQN":
            # 用训练网络选择动作
            max_action = self.q_net(next_states).max(1)[1].view(-1, 1)
            # 用目标网络计算Q值
            max_next_q_values = self.target_q_net(next_states).gather(1, max_action)
        else:  # DQN和DuelingDQN更新方式一样
            # 下个状态的最大Q值 Q(s',a')
            max_next_q_values = self.target_q_net(next_states).max(1)[0].view(-1, 1)
        # TD error target  r+gamma*Q(s',a')
        q_targets = rewards + self.gamma * max_next_q_values * (1 - dones)

        dqn_loss = torch.mean(F.mse_loss(q_values, q_targets))

        # 网络参数更新
        self.optimizer.zero_grad()
        dqn_loss.backward()
        self.optimizer.step()

        if self.count % self.target_update == 0:
            # 更新目标网络
            self.target_q_net.load_state_dict(self.q_net.state_dict())
        self.count += 1




if __name__ == '__main__':

    lr = 2e-3
    num_episodes = 500
    hidden_dim = 128
    gamma = 0.98
    epsilon = 0.01
    # 每10次更新，更新一次目标网络
    target_update = 10
    buffer_size = 10000
    minimal_size = 500
    batch_size = 64
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device(
        "cpu")

    env_name = 'CartPole-v0'
    env = gym.make(env_name)
    random.seed(0)
    np.random.seed(0)
    env.seed(0)
    torch.manual_seed(0)
    replay_buffer = ReplayBuffer(buffer_size)
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n
    agent = DQN(state_dim, hidden_dim, action_dim, lr, gamma, epsilon,
                target_update, device)

    return_list = []
    for i in range(10):
        with tqdm(total=int(num_episodes / 10), desc='Iteration %d' % i) as pbar:
            for i_episode in range(int(num_episodes / 10)):
                episode_return = 0
                state = env.reset()
                done = False
                while not done:
                    action = agent.take_action(state)
                    next_state, reward, done, _ = env.step(action)
                    replay_buffer.add(state, action, reward, next_state, done)
                    state = next_state
                    episode_return += reward
                    # 当buffer数据的数量超过一定值后,才进行Q网络训练
                    if replay_buffer.size() > minimal_size:
                        b_s, b_a, b_r, b_ns, b_d = replay_buffer.sample(batch_size)
                        transition_dict = {
                            'states': b_s,
                            'actions': b_a,
                            'next_states': b_ns,
                            'rewards': b_r,
                            'dones': b_d
                        }
                        agent.update(transition_dict)
                return_list.append(episode_return)
                if (i_episode + 1) % 10 == 0:
                    pbar.set_postfix({
                        'episode':
                            '%d' % (num_episodes / 10 * i + i_episode + 1),
                        'return':
                            '%.3f' % np.mean(return_list[-10:])
                    })
                pbar.update(1)
    episodes_list = list(range(len(return_list)))
    plt.plot(episodes_list, return_list)
    plt.xlabel('Episodes')
    plt.ylabel('Returns')
    plt.title('DQN on {}'.format(env_name))
    plt.show()

    mv_return = utils.moving_average(return_list, 9)
    plt.plot(episodes_list, mv_return)
    plt.xlabel('Episodes')
    plt.ylabel('Returns')
    plt.title('DQN on {}'.format(env_name))
    plt.show()
